#!/usr/bin/env python3
#coding:utf-8

__author__ = 'xmxoxo<xmxoxo@qq.com>'

'''
mq_lib.py
类库，包括MQ调用，RPC模式等
'''

import pika
import redis
import time
import random
import json
import logging
from config import *

# RabbitMQ类
class MQ():
    def __init__ (self, ip='127.0.0.1', port=5672):
        # MQ服务器IP和端口
        self.ip = ip
        self.port = port
        # 默认自动重连
        self.autoconnect = True
        self.conn = None
        self.channel = None
        self.queue_name = ""
        self.dofun = None

    # 用户登录
    def login(self, user, pwd):
        #建立连接
        userx = pika.PlainCredentials(user, pwd)
        parm = pika.ConnectionParameters(self.ip, self.port, '/', credentials=userx)
        self.conn = pika.BlockingConnection(parm)
        print('已连接到MQ:')
        return 1

    # 发送消息
    def send(self, channel, msg):
        pass
        # 创建队列，持久化，topic
        if self.channel is None:
            #开辟管道
            channelx = self.conn.channel()
            #声明队列，参数为队列名
            #channelx.queue_declare(queue=channel)
            channelx.exchange_declare(exchange='topic_exchange', exchange_type='topic', durable=True)
            self.channel = channelx

        # 发送数据，发送一条，如果要发送多条则复制此段
        self.channel.basic_publish(exchange="topic_exchange",
                               routing_key=channel, # 队列名
                               body=msg  # 发送的数据
                               )
        print("--------发送数据完成-----------")
        return 1

    def close(self):
        # 关闭连接
        self.conn.close()

    # 绑定频道接收事件, callback
    def bind(self, channel, dofun=None, exclusive_consumption=False):
        ''' 绑定事件
        exclusive_consumption：是否独占
        '''
        #connection = pika.BlockingConnection(pika.ConnectionParameters(host='localhost'))
        #channelx = connection.channel()
        if not dofun is None:
            self.dofun = dofun

        if self.channel is None:
            # 开辟管道
            channelx = self.conn.channel()
            channelx.exchange_declare(exchange='topic_exchange', exchange_type='topic', durable=True)
            #channelx.exchange_declare(exchange='topic_logs', exchange_type='topic', durable=True)

            # 不指定queue名字,rabbit会随机分配一个名字,exclusive=True会在使用此queue的消费者断开后,自动将queue删除
            #result = channelx.queue_declare(exclusive=True)
            # 获取随机的queue名字
            #queue_name = result.method.queue
            #print("random queuename:", queue_name)

            # 定义队列
            self.queue_name = "queue_%s" % str(int(time.time()))
            # channelx.queue_declare(queue=self.queue_name, durable=True, auto_delete=True)
            
            if exclusive_consumption:
                # 创建独占队列
                result = channelx.queue_declare(queue=self.queue_name, durable=False, exclusive=True)
            else:
                # 创建非独占、持久化的队列
                result = channelx.queue_declare(queue=self.queue_name, durable=True, exclusive=False)

            self.channel = channelx

        # queue绑定到转发器上
        if type(channel)==str:
            self.channel.queue_bind(exchange='topic_exchange', queue=self.queue_name, routing_key=channel)
            #self.channel.queue_bind(exchange='topic_logs', queue=self.queue_name, routing_key=channel)

        # channel 支持多个；如果是列表则绑定多个
        # 使用“|”符号自动拆分 2023/10/7
        if '|' in channel:
            channel_list = filter(None, channel.split('|'))
            for key in channel_list:
                print('正在绑定:%s'% key)
                self.channel.queue_bind(exchange='topic_exchange', queue=self.queue_name, routing_key=key)
                
        # 指定QOS
        # self.channel.basic_qos(prefetch_count=1)

        # 消费消息 consume_callback为回调函数
        self.channel.basic_consume( 
            on_message_callback=self.callback,  # 如果收到消息，就调用callback函数来处理消息
            queue=self.queue_name,              # 你要从哪个队列里收消息
            # no_ack=True                       # 写的话，如果接收消息，机器宕机消息就丢了; 一般不写。宕机则生产者检测到发给其他消费者
            # auto_ack=True
            )
         
        print(' [*] Waiting for messages. To exit press CTRL+C')
        self.channel.start_consuming() # 开始消费消息


        return 1

    def callback(self, channel, method, properties, body): # 四个参数为标准格式
        # print(channel, method, properties) # 打印看一下是什么
        # print('来自:', method.routing_key)
        # 管道内存对象 内容相关信息 
        # Todo: 处理数据
        #print("[%s]收到消息: %r" % (method.routing_key, body) )
        #time.sleep(1)
        try:
            #print('self.dofun:', self.dofun)
            if not self.dofun is None:
                #print('do fun...')
                # 组织数据
                ret_dat = {'channel':method.routing_key, 'body':body}
                self.dofun(ret_dat)
                channel.basic_ack(delivery_tag=method.delivery_tag) # 告诉生成者，消息处理完成
            else:
                print('No callback')
        except Exception as e:
            print(e)
        
# -----------------------------------------

# RPC客户端 
class RPC_client():
    def __init__ (self, mq):
        self.mq = mq

    # await
    def send(self, msg):
        pass

# RPC服务端 
class RPC_svr():
    def __init__ (self, mq):
        pass
            
    def bind (self, channel):
        pass

# 消息类
class MSG():
    def __init__ (self):
        self.data = {}

    # 设置数据直接解析
    def setdat (self, dat:str):
        self.data = json.loads(dat)

    def update (self, dat:str):
        if type(dat) == str:
            mdat = json.loads(dat)
            self.data.update(mdat)
        if type(dat) == dict:
            self.data.update(dat)
        
    def add (self, key, value):
        ''' 增加数据'''

        self.data[key] = value
    
    # 返回字符串
    def __str__(self):
        return json.dumps(self.data)

    def rand_id(self, key:str, prekey=''):
        ''' 生成随机的ID 
        '''
        if key == '':
            return ''

        msgid = prekey + str(int(time.time()*1000)) + str(random.randrange(100, 999))
        self.data[key] = msgid
        return msgid

# -----------------------------------------

class OBJREDIS():
    ''' Redis 读写数据
    '''
    def __init__(self, host='127.0.0.1', port=6379, dat_table="DAT_TABLE"):
        self.host = host
        self.port = port
        self.pool = None
        self.r = None

        # 表名称
        self.dat_table = dat_table
        
        self.conn_redis()

    def conn_redis(self):
        try:
            pool = redis.ConnectionPool(host=self.host, port=self.port,
                                            max_connections=50, decode_responses=True)
            self.r = redis.Redis(connection_pool=pool)
            logging.info('redis inited...')
        except Exception as e:
            print(e)
            pass

        
    def set_record(self, datid:str, value):
        '''写入记录
        '''
        #dat = json.dumps(value)
        #r = self.r.hset(self.dat_table, datid, dat)
        key = '%s_%s' % (self.dat_table, datid)
        r = self.r.set(key, value)
        return r

    def get_record(self, datid:str):
        '''读取记录
        '''
        ret = {}
        #dat = self.r.hget(self.dat_table, datid)
        key = '%s_%s' % (self.dat_table, datid)
        dat = self.r.get(key)
        if dat:
            #ret = json.loads(dat)
            ret = bytes.decode(dat)
        else:
            ret = None
        return ret
    
    def lock(self, key, value, time_out=30):
        ''' 锁定，成功返回1
        time_out: 过期时间，单位秒
        '''
        ret = self.r.setnx(key, value)
        if ret:
            self.r.expire(key, time_out)
            return 1
        else:
            return 0
    
    def lock1(self, key, value):
        pipe = self.r.pipeline()

        while True:
            try:
                pipe.watch(key)  # 监视键名为 key 的键
                value = int(pipe.get(key))  # 获取 key 的值
                value = value + 1
                pipe.multi()  # 开启事务
                pipe.setnx(key, value)
                pipe.expire(key, 60) #  过期时间，单位秒
                pipe.execute()  # 执行事务
                break
            except redis.WatchError:
                continue

        result = r.get(msgid)
        print(result)
        return 

    # 尝试获取分布式锁
    def get_lock(self, lock_key, timeout=30):
        ''' timeout为操作的超时时间，单位秒
        '''

        # 设置锁的超时时间
        expire_time = time.time() + timeout

        while True:
            # 尝试获取锁
            if self.r.setnx(lock_key, 1):
                # 获取锁成功，设置锁的超时时间
                self.r.expire(lock_key, timeout)
                return True

            # 获取锁失败，检查锁是否已经超时
            current_time = time.time()
            if current_time > expire_time:
                # 锁超时，返回获取锁失败
                return False

            # 休眠一段时间后重新尝试获取锁
            time.sleep(0.1)

def test_redis():
    ''' 单元测试
    '''
    print('正在创建Redis数据库连接...')
    redis_table_name = 'MESSAGE_REPEAT'
    obj_redis = OBJREDIS(dat_table=redis_table_name)

    print('正在写入数据...')

    mid = str(int(time.time()*1000))
    obj_redis.set_record(mid, 0)

    # 读取：
    print('正在读取数据...')
    nvalue = obj_redis.get_record(mid)
    print("return value:", nvalue)


if __name__ == '__main__':
    pass

    # test_redis()


